#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Safe code execution utilities for the visualization engine
"""

import pandas as pd
import numpy as np
from typing import Dict, Any, Tuple, Optional
import time
import traceback
from io import StringIO
from contextlib import redirect_stdout, redirect_stderr
from data_engine.utils.validators import CodeValidator


class SafeCodeExecutor:
    """安全代码执行器"""
    
    def __init__(self, timeout: int = 30):
        self.timeout = timeout
        self._setup_safe_globals()
    
    def _setup_safe_globals(self) -> None:
        """设置安全的全局变量环境"""
        # 允许的内置函数
        safe_builtins = {
            # 基本类型和函数
            'len': len, 'str': str, 'int': int, 'float': float, 'bool': bool,
            'list': list, 'dict': dict, 'tuple': tuple, 'set': set,
            'range': range, 'enumerate': enumerate, 'zip': zip,
            'sum': sum, 'min': min, 'max': max, 'abs': abs, 'round': round,
            'sorted': sorted, 'reversed': reversed, 'any': any, 'all': all,
            'isinstance': isinstance, 'type': type, 'hasattr': hasattr,
            
            # 异常类
            'Exception': Exception, 'ValueError': ValueError, 'TypeError': TypeError,
            'KeyError': KeyError, 'IndexError': IndexError, 'AttributeError': AttributeError,
            
            # 必要的内置函数（支持import等）
            '__import__': __import__,
            '__builtins__': __builtins__,
        }
        
        self.safe_globals = {
            '__builtins__': safe_builtins,
            'pd': pd, 'pandas': pd, 'np': np, 'numpy': np,
        }
    
    def execute_code(self, code: str, df: pd.DataFrame) -> Tuple[bool, pd.DataFrame, str, float]:
        """安全执行数据处理代码"""
        start_time = time.time()
        
        # 验证代码安全性
        is_safe, warnings = CodeValidator.validate_processing_code(code)
        if not is_safe:
            print(f"代码安全检查发现问题: {warnings}")
        
        # 验证代码语法
        is_valid, syntax_error = CodeValidator.validate_code_syntax(code)
        if not is_valid:
            execution_time = time.time() - start_time
            return False, df, f"代码语法错误: {syntax_error}", execution_time
        
        try:
            # 设置执行环境
            local_vars = self.safe_globals.copy()
            local_vars['df'] = df.copy()
            
            # 捕获输出
            stdout_capture = StringIO()
            stderr_capture = StringIO()
            
            with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture):
                exec(code, local_vars, local_vars)
            
            # 获取结果DataFrame
            result_df = local_vars.get('df', df)
            
            if not isinstance(result_df, pd.DataFrame):
                raise ValueError("代码执行后df不是DataFrame类型")
            
            execution_time = time.time() - start_time
            return True, result_df, "", execution_time
            
        except Exception as e:
            execution_time = time.time() - start_time
            error_message = f"代码执行失败: {str(e)}"
            return False, df, error_message, execution_time


class DataProcessingTemplates:
    """数据处理代码模板"""
    
    @staticmethod
    def get_aggregation_template(group_by: str, agg_field: str, agg_func: str) -> str:
        return f"df = df.groupby('{group_by}')['{agg_field}'].{agg_func}().reset_index()"
    
    @staticmethod
    def get_filter_template(condition: str) -> str:
        return f"df = df.query('{condition}')"
    
    @staticmethod
    def get_sort_template(sort_by: str, ascending: bool = True) -> str:
        return f"df = df.sort_values('{sort_by}', ascending={ascending})"
    
    @staticmethod
    def get_limit_template(limit: int) -> str:
        return f"df = df.head({limit})"